using System;
using System.Collections;
using System.Text;

using Mono.Cecil;
using Mono.Cecil.Cil;

using CILUtil;

namespace CILFlow
{
    public class CILFlowGraph
    {
        CILInstructionBlock[] m_blocks;
        public CILInstructionBlock[] Blocks
        {
            get { return m_blocks;  }
        }
        public CILFlowGraph(CILInstructionBlock[] blocks)
        {
            this.m_blocks = blocks;
        }
        private string Quote(object obj)
        {
            return "\"" + obj + "\"";
        }
        public string ToDot()
        {
            StringBuilder buffer = new StringBuilder();
            buffer.AppendLine( "digraph G {" );
            foreach( CILInstructionBlock block in this.m_blocks )
            {
                if (block.Children.Count == 0)
                {
                    buffer.AppendLine( Quote(block) + ";" );
                }
                foreach (CILInstructionBlock next in block.Children)
                {
                    buffer.AppendLine( Quote(block) + "->" + Quote(next) + ";");
                }
            }
            buffer.AppendLine( "}" );
            return buffer.ToString();
        }

        public string FullGraphToDot( MethodDefinition md)
        {
            StringBuilder buffer = new StringBuilder();
            buffer.AppendLine("digraph G {");
            foreach (CILInstructionBlock block in this.m_blocks)
            {
                InstructionsSubGraph(block, md, buffer);

                //if (block.Children.Count == 0)
                //{
                //    buffer.AppendLine("cluster_" + block.GetHashCode() + ";");
                //}
                foreach (CILInstructionBlock next in block.Children)
                {
                    buffer.AppendLine(block.End.GetHashCode() + " -> " + next.Start.GetHashCode() + ";");
                }

                
            }
            buffer.AppendLine("}");
            return buffer.ToString();
        }

        public void InstructionsSubGraph(CILInstructionBlock block, MethodDefinition md, StringBuilder buffer)
        {
            PopTreeList list = PopTreeListBuilder.Construct(md, block);
            buffer.Append("subgraph cluster_").Append(block.GetHashCode()).AppendLine("{");
            if( list.PopTrees.Count == 0) 
            {
                buffer.AppendLine(Quote(block));
            }
            foreach (PopTree tree in list.PopTrees)
            {
                tree.ToDot(md,buffer);
            }
            buffer.AppendLine("}");
        }
    }

    public class CILInstructionBlock
    {
        public CILInstructionBlock(ArrayList instructions)
        {
            m_instructions = new Instruction[instructions.Count];
            instructions.CopyTo( m_instructions );
        }
        private Instruction[] m_instructions;
        private bool m_isCatchBlock;
        public bool IsCatchBlock
        {
            get { return m_isCatchBlock;  }
            internal set { m_isCatchBlock = value; }
        }
        public Instruction[] Instructions
        {
            get { return m_instructions; }
        }
        public Instruction Start
        {
            get { return m_instructions[0]; }
        }
        public Instruction End
        {
            get { return m_instructions[m_instructions.Length-1]; }
        }
        private ArrayList m_children = new ArrayList();
        public ArrayList Children
        {
            get { return m_children; }
        }
        private string Escape( string str )
        {
            return str.Replace( "\"","\\"+"\"");
        }
        public override string ToString()
        {
            return Escape(CecilFormatter.FormatInstruction(Start)) + "\\n" +
                   Escape(CecilFormatter.FormatInstruction(End));
        }
    }

    public class CILMethodBlockConstructor
    {
        private int[] m_leaders;
        private int[] m_instructionBlockMap;
        private int[] m_IsExceptionHandlerStart;
        private Hashtable m_instructionIndexMap;

        private CILMethodBlockConstructor(MethodDefinition md)
        {
            md.Body.Simplify();

            m_leaders = new int[md.Body.Instructions.Count];
            m_instructionBlockMap = new int[md.Body.Instructions.Count];
            m_IsExceptionHandlerStart = new int[md.Body.Instructions.Count];

            m_instructionIndexMap = new Hashtable();
            int idx = 0;
            foreach (Instruction instruction in md.Body.Instructions)
            {
                m_instructionBlockMap[idx] = -1;
                m_instructionIndexMap[instruction] = idx;
                idx++;
            }
        }

        public static CILFlowGraph ConstructFlow(MethodDefinition md)
        {
            CILMethodBlockConstructor instance = new CILMethodBlockConstructor(md);
            instance.MarkBlockStarts(md);
            CILInstructionBlock[] blocks = instance.BuildBlocks(md);
            instance.ConnectBlocks(blocks);
            return new CILFlowGraph( blocks );
        }

        private void MarkLeader(Instruction i)
        {
            int idx = (int)m_instructionIndexMap[i];
            m_leaders[idx] = 1;
        }
        private void MarkIsExceptionHandlerStart(Instruction i)
        {
            int idx = (int)m_instructionIndexMap[i];
            m_IsExceptionHandlerStart[idx] = 1;
        }

        private bool IsExceptionStart(Instruction i)
        {
            int idx = (int)m_instructionIndexMap[i];
            return m_IsExceptionHandlerStart[idx] == 1;
        }

        private bool IsLeader(Instruction i)
        {
            int idx = (int)m_instructionIndexMap[i];
            return m_leaders[idx] == 1;
        }

        private void MarkBlockStarts(MethodDefinition md)
        {
            Instruction instruction = md.Body.Instructions[0];

            // the first instruction starts a block
            MarkLeader(instruction);
            for (int i = 1; i < md.Body.Instructions.Count; ++i)
            {
                instruction = md.Body.Instructions[i];
                if (CodeProperties.IsBlockEnd(instruction))
                {
                    MarkLeadersFromBlockEnd(instruction);
                }
            }
            // Handle exceptions handlers.
            MarkExceptions(md);
        }

        private void MarkExceptions(MethodDefinition md)
        {
            foreach (ExceptionHandler exp in md.Body.ExceptionHandlers)
            {   
                // The end of a try block is a branch.  (Will mark next, start of catch)
                if (exp.TryEnd != null)
                {
                    MarkLeadersFromBlockEnd(exp.TryEnd);
                }
                // The end of the handler is a branch.  (Will mark next, start of another catch,finally).
                if (exp.HandlerEnd != null)
                {
                    MarkLeadersFromBlockEnd(exp.HandlerEnd);
                    // Want to know if this block will be the start of catch block.
                    // Therefore, we might want to simulate that the exception is being pushed onto the stack
                    // with PopTrees.
                    MarkIsExceptionHandlerStart(exp.HandlerStart);
                }

                if (exp.FilterEnd != null)
                {
                    MarkLeadersFromBlockEnd(exp.FilterEnd);
                    MarkIsExceptionHandlerStart(exp.FilterStart);
                }
            }
        }

        private void MarkLeadersFromBlockEnd(Instruction instruction)
        {
            // the target of a branch starts a block
            Instruction[] targets = Targets(instruction);
            if (targets != null)
            {
                foreach (Instruction target in targets)
                {
                    MarkLeader(target);
                }
            }
            // the next instruction after a branch starts a block
            if (null != instruction.Next) MarkLeader(instruction.Next);
        }

        private void ConnectBlocks(CILInstructionBlock[] cilBlocks)
        {
            foreach (CILInstructionBlock block in cilBlocks)
            {
                if (CodeProperties.IsBlockEnd(block.End))
                {
                    Instruction[] targets = Targets(block.End);
                    if (targets != null)
                    {
                        foreach (Instruction target in targets)
                        {
                            block.Children.Add(BlockFromInstruction(target, cilBlocks));
                        }
                    }
                    // For conditional branches, add the next instruction as a child.
                    if (block.End.OpCode.FlowControl == FlowControl.Cond_Branch && block.End.Next != null)
                    {
                        block.Children.Add(BlockFromInstruction( block.End.Next, cilBlocks ));
                    }
                }
                // Sometimes flow control can be a call.  Just always connect to next guy.
                else if (/*block.End.OpCode.FlowControl == FlowControl.Next &&*/ block.End.Next != null)
                {
                    block.Children.Add(BlockFromInstruction(block.End.Next, cilBlocks));
                }

            }
        }

        private void UpdateBlockInstructionMap(CILInstructionBlock block, int blockNum)
        {
            foreach (Instruction i in block.Instructions)
            {
                int idx = (int)m_instructionIndexMap[i];
                if (m_instructionBlockMap[idx] != -1)
                    throw new Exception("Overlapping Block Assignment Detected");
                m_instructionBlockMap[idx] = blockNum;
            }
        }

        private CILInstructionBlock BlockFromInstruction(Instruction i, CILInstructionBlock[] blocks)
        {
            int idx = (int)m_instructionIndexMap[i];
            int bindex = (int)m_instructionBlockMap[idx];
            return blocks[bindex];
        }

        private CILInstructionBlock[] BuildBlocks( MethodDefinition md )
        {
            ArrayList blocks = new ArrayList();
            for(int pc=0; pc<md.Body.Instructions.Count; pc++)
            {
                Instruction instruction = md.Body.Instructions[pc];
                if (IsLeader(instruction))
                {
                    CILInstructionBlock block = ConstructBlock(instruction, md);
                    UpdateBlockInstructionMap(block, blocks.Count);
                    blocks.Add(block);
                }
            }
            CILInstructionBlock[] cilBlocks = new CILInstructionBlock[blocks.Count];
            blocks.CopyTo(cilBlocks);
            return cilBlocks; 
        }

        private CILInstructionBlock ConstructBlock(Instruction i, MethodDefinition md)
        {
            ArrayList instructions = new ArrayList();
            instructions.Add(i);
            int pc = (int)m_instructionIndexMap[i] + 1;
            bool isCatchStart = false;
            while( pc < md.Body.Instructions.Count && !IsLeader( md.Body.Instructions[pc] ) )
            {
                Instruction cur = md.Body.Instructions[pc];
                instructions.Add(cur);
                if (isCatchStart == false && IsExceptionStart(cur))
                    isCatchStart = true;
                pc++;
            }
            CILInstructionBlock block = new CILInstructionBlock(instructions);
            block.IsCatchBlock = isCatchStart;
            return block;
        }

        private Instruction[] Targets(Instruction instruction)
        {
            Instruction target = instruction.Operand as Instruction;
            if (target == null)
            {
                return (instruction.Operand as Instruction[]);
            }
            return new Instruction[]{ target };
        }
    }
}
